import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import models

class VGGFT(nn.Module):
    """
    VGG微调网络实现，基于torchvision的VGG16模型
    对应于train_vgg_ft.prototxt的Caffe模型
    """
    def __init__(self, pretrained=False):
        super(VGGFT, self).__init__()
        # 加载预训练的VGG16模型
        # 使用新的weights参数代替pretrained
        if pretrained:
            vgg16 = models.vgg16(weights=models.VGG16_Weights.IMAGENET1K_V1)
        else:
            vgg16 = models.vgg16(weights=None)
        
        # 提取特征部分 (卷积层)
        self.features = vgg16.features

        # avgpool
        self.avgpool = nn.AdaptiveAvgPool2d((7, 7))
        
        # 提取分类器部分 (全连接层)，但不包括最后一层
        # 对应于Caffe中的fc6_m和fc7_m
        self.fc6_m = vgg16.classifier[0]
        self.relu6 = vgg16.classifier[1]
        self.drop6 = vgg16.classifier[2]
        self.fc7_m = vgg16.classifier[3]
        self.relu7 = vgg16.classifier[4]
        self.drop7 = vgg16.classifier[5]
        
        # 添加我们自己的输出层，对应于Caffe中的fc8
        self.fc8 = nn.Linear(4096, 1)
        
        # 初始化权重
        self._initialize_weights()
        
    def forward(self, x):
        # 特征提取
        x = self.features(x)

        # avgpool
        x = self.avgpool(x)
        
        # 展平特征图
        x = x.view(x.size(0), -1)
        
        # 全连接层
        x = self.fc6_m(x)
        x = self.relu6(x)
        x = self.drop6(x)
        x = self.fc7_m(x)
        x = self.relu7(x)
        x = self.drop7(x)
        
        # 最终输出层
        x = self.fc8(x)
        
        return x
    
    def _initialize_weights(self):
        """初始化fc6_m, fc7_m和fc8的权重"""
        for m in [self.fc6_m, self.fc7_m, self.fc8]:
            nn.init.xavier_uniform_(m.weight)
            nn.init.constant_(m.bias, 0)
    
    def get_parameter_groups(self):
        """
        获取参数组，模拟Caffe中的lr_mult和decay_mult
        
        在Caffe中：
        - 卷积层权重的lr_mult=1, decay_mult=1
        - 卷积层偏置的lr_mult=2, decay_mult=0
        - fc6_m/fc7_m/fc8层的lr_mult=10(权重)/20(偏置), decay_mult=1/0
        
        基础学习率为0.00001（来自solver_vgg.prototxt）
        
        返回:
            参数组列表，可直接用于优化器
        """
        # 基础学习率
        base_lr = 0.00001
        
        # 基础层参数 (卷积层)
        base_params = []
        # fc6_m, fc7_m, fc8层参数
        fc_params = []
        
        # 遍历所有命名参数
        for name, param in self.named_parameters():
            # 根据参数名称判断是否为fc6_m, fc7_m或fc8层
            if 'fc6_m' in name or 'fc7_m' in name or 'fc8' in name:
                if 'weight' in name:
                    # fc权重 - lr_mult=10, decay_mult=1
                    fc_params.append({'params': param, 'lr': base_lr * 10, 'weight_decay': 0.0005})
                else:
                    # fc偏置 - lr_mult=20, decay_mult=0
                    fc_params.append({'params': param, 'lr': base_lr * 20, 'weight_decay': 0})
            else:
                if 'weight' in name:
                    # 卷积层权重 - lr_mult=1, decay_mult=1
                    base_params.append({'params': param, 'lr': base_lr * 1, 'weight_decay': 0.0005})
                else:
                    # 卷积层偏置 - lr_mult=2, decay_mult=0
                    base_params.append({'params': param, 'lr': base_lr * 2, 'weight_decay': 0})
        
        # 返回两组参数
        return base_params + fc_params
    
    def load_from_rank_model(self, rank_model_path):
        """
        从预训练的排序模型加载权重
        
        Args:
            rank_model_path: 预训练排序模型的路径
        """
        # 加载预训练模型
        rank_model_dict = torch.load(rank_model_path)
        
        # 创建一个新的状态字典，只包含我们需要的层
        model_dict = self.state_dict()
        
        # 过滤出需要加载的参数（排除fc8层，因为输出可能不同）
        # 注意：这里假设预训练模型和当前模型的层名称匹配
        # 如果层名称不匹配，需要手动映射
        pretrained_dict = {k: v for k, v in rank_model_dict.items() 
                          if k in model_dict and 'fc8' not in k}
        
        # 更新当前模型的参数
        model_dict.update(pretrained_dict)
        
        # 加载参数
        self.load_state_dict(model_dict)
        
        print(f"成功从 {rank_model_path} 加载预训练权重")